import java.text.DecimalFormat;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.text.*;
import java.util.*;
import java.io.*;

public class SecondOrderHMM extends HMM {

  // If this parameter is set to true, we use the smoothing technique for state transitions described in http://acl.ldc.upenn.edu/P/P99/P99-1023.pdf
  private static boolean fancySmoothing = true;

  double[][][] loga;
     
  private static double[][] toFirstOrder ( double[][][] amat ) {
	double[][] aux = new double[amat.length][amat.length];
	for ( int i = 0 ; i < amat.length; i++ ) for ( int j = 0 ; j < amat.length ; j++ ) aux[i][j] = 0.0;
	for ( int i = 0 ; i < amat.length; i++ ) for ( int j = 0 ; j < amat.length ; j++ ) for ( int k = 0 ; k < amat.length ; k++ ) aux[j][k] += amat[i][j][k];
	return aux;
  }

  public SecondOrderHMM(Object[] state, double[][][] amat, Object[] esym, double[][] emat) {
    super(state,toFirstOrder(amat),esym,emat);
    if (state.length != amat.length) throw new IllegalArgumentException("HMM: state and amat disagree");
    if (amat.length != emat.length) throw new IllegalArgumentException("HMM: amat and emat disagree");
    for (int i=0; i<amat.length; i++) {
      if (state.length != amat[i].length) throw new IllegalArgumentException("HMM: amat non-square");
      if (esym.length != emat[i].length) throw new IllegalArgumentException("HMM: esym and emat disagree");
    }
    nstate = state.length + 1;
    this.state = new Object[nstate];
    loga = new double[nstate][nstate][nstate];
    for (int i=0; i<nstate; i++) for (int j=0; j<nstate; j++) for (int k=0; k<nstate; k++) loga[i][j][k] = 0.0;
    this.state[0] = (Object)(new String("---INITIAL-STATE---"));
    loga[0][0][0] = Double.NEGATIVE_INFINITY;
    double fromstart = Math.log(1.0 / state.length);
    for (int j=1; j<nstate; j++) loga[0][0][j] = fromstart;
    for (int i=1; i<nstate; i++) {
      this.state[i] = state[i-1];
      loga[i][0][0] = Double.NEGATIVE_INFINITY;
      loga[0][i][0] = Double.NEGATIVE_INFINITY;
      for (int j=1; j<nstate; j++) {
       loga[i][j][0] = Double.NEGATIVE_INFINITY;
       loga[i][0][j] = Double.NEGATIVE_INFINITY;
       for (int k=1; k<nstate; k++) {
		loga[i][j][k] = Math.log(amat[i-1][j-1][k-1]);
	        loga[0][j][k] += amat[i-1][j-1][k-1];
       }
      }
    }
    for (int j=1; j<nstate; j++) for (int k=1; k<nstate; k++) loga[0][j][k] = Math.log(loga[0][j][k]);
    this.esym = esym;
    nesym = esym.length;
    loge = new double[emat.length + 1][esym.length];
    for (int b=0; b < nesym; b++) {
      loge[0][b] = Double.NEGATIVE_INFINITY;
      for (int k=0; k<emat.length; k++) loge[k+1][b] = Math.log(emat[k][b]);
    }    
  }

  @Override
  public void printa(PrintStream out) {
    out.println("Transition probabilities:");
    for (int i=1; i<nstate; i++) {
      for (int j=1; j<nstate; j++) for (int k=1; k<nstate; k++) out.print(fmtlog(loga[j][k][i]));
      out.println();
    }
  }

  public static SecondOrderHMM baumwelch(Object[][] xs, Object[] state, Object[] esym, final double threshold ) {
    int nstate = state.length;
    int nseqs  = xs.length;
    int nesym  = esym.length;
    SecondOrderForward[] fwds = new SecondOrderForward[nseqs];
    SecondOrderBackward[] bwds = new SecondOrderBackward[nseqs];
    double[] logP = new double[nseqs];
    double[][][] amat = new double[nstate][nstate][];
    double[][] emat = new double[nstate][];
    int[] esyminv = new int[esym.length];
    for (int i=0; i<esyminv.length; i++) esyminv[i] = -1;
    for (int b=0; b<nesym; b++) esyminv[b] = b;
    for (int k=0; k<nstate; k++) for (int l=0; l<nstate; l++) amat[k][l] = uniformdiscrete(nstate);
    for (int k=0; k<nstate; k++) emat[k] = randomdiscrete(nesym);
    SecondOrderHMM hmm = new SecondOrderHMM(state, amat, esym, emat);  
    double oldloglikelihood;
    double loglikelihood = fwdbwd(hmm, xs, fwds, bwds, logP);
    int numIterations = 0;
    do {
      if ( maxIterationsBaumWelsh > -1 && ( numIterations++ > maxIterationsBaumWelsh ) ) break;
      oldloglikelihood = loglikelihood;
      double[][][] A = new double[nstate][nstate][nstate];
      double[][] E = new double[nstate][nesym];
      for (int s=0; s < nseqs; s++) {
        Object[] x = xs[s];
        SecondOrderForward fwd  = (SecondOrderForward)(fwds[s]);
        SecondOrderBackward bwd = (SecondOrderBackward)(bwds[s]);
        int L = x.length;
        double P = logP[s];
        for (int i=0; i < L; i++) {
          for (int k=0; k<nstate; k++) E[k][esyminv[hmm.getSym(x[i])]] += exp(fwd.f[i+1][k+1] + bwd.b[i+1][k+1] - P);
        }
        for (int i=0; i<L-1; i++) 
          for (int j=0; j<nstate; j++) 
	    for (int k=0; k<nstate; k++) 
	      for (int ell=0; ell<nstate; ell++) 
	              A[j][k][ell] += exp(fwd.f[i+1][k+1] 
	                           + hmm.loga[j+1][k+1][ell+1] 
	                           + hmm.loge[ell+1][hmm.getSym(x[i+1])] 
	                           + bwd.b[i+2][ell+1] 
	                           - P);
      }
      if ( fancySmoothing ) {
	double c0 = 0;
        for (int j=0; j<nstate; j++) for (int k=0; k<nstate; k++) for (int ell=0; ell<nstate; ell++) c0 += A[j][k][ell];
        for (int j=0; j<nstate; j++) {
           for (int k=0; k<nstate; k++) {
             double c1 = 0.0;
	     for (int jj=0; jj<nstate; jj++) for (int ell=0; ell<nstate; ell++) c1 += A[jj][k][ell];
	     double c2 = 0.0;
	     for (int ell=0; ell<nstate; ell++) c2 += A[j][k][ell];
             for (int ell=0; ell<nstate; ell++) {
		double n1 = 0.0;
                for (int jj=0; jj<nstate; jj++) for (int kk=0; kk<nstate; kk++) n1 += A[jj][kk][ell];
		double n2 = 0.0;
                for (int jj=0; jj<nstate; jj++) n2 += A[jj][k][ell];
		double n3 = A[j][k][ell];
                double k2 = (Math.log(n2 + 1.0) + 1.0) / (Math.log(n2 + 1.0) + 2.0);
                double k3 = (Math.log(n3 + 1.0) + 1.0) / (Math.log(n3 + 1.0) + 2.0);
                amat[j][k][ell] = (k3*(n3/c2)) + (((1.0-k3)*k2)*(n2/c1)) + (((1-k3)*(1-k2))*(n1/c0));
             }
           }
        }
        for (int j=0; j<nstate; j++) for (int k=0; k<nstate; k++) {
          double Aksum = 0;
          for (int ell=0; ell<nstate; ell++) Aksum += amat[j][k][ell];
          for (int ell=0; ell<nstate; ell++) amat[j][k][ell] = amat[j][k][ell] / Aksum;
        }
      } else for (int j=0; j<nstate; j++) for (int k=0; k<nstate; k++) {
        double Aksum = 0;
        for (int ell=0; ell<nstate; ell++) Aksum += A[j][k][ell];
        for (int ell=0; ell<nstate; ell++) amat[j][k][ell] = ( A[j][k][ell] + 1.0 )  / ( Aksum + 0.0 + ( nstate * nstate) );
      }
      for (int k=0; k<nstate; k++) {
        double Eksum = 0;
        for (int b=0; b<nesym; b++) Eksum += E[k][b];
        for (int b=0; b<nesym; b++) emat[k][b] = ( E[k][b] + 1.0 ) / ( Eksum + 0.0 + nesym);
      }
      hmm = new SecondOrderHMM(state, amat, esym, emat);
      loglikelihood = fwdbwd(hmm, xs, fwds, bwds, logP);
      if ( debugBaumWelsh ) System.err.println("Baum-Welch Training - log likelihood = " + loglikelihood + " ( diff from previous is " + Math.abs(oldloglikelihood - loglikelihood) + ")");
    } while (Math.abs(oldloglikelihood - loglikelihood) > threshold);   
    return hmm;
  }

  protected static double fwdbwd(HMM hmm, Object[][] xs, Forward[] fwds, Backward[] bwds, double[] logP) {
    double loglikelihood = 0;
    for (int s=0; s<xs.length; s++) {
      fwds[s] = new SecondOrderForward((SecondOrderHMM)hmm, xs[s]);
      bwds[s] = new SecondOrderBackward((SecondOrderHMM)hmm, xs[s]);
      logP[s] = ((SecondOrderForward)(fwds[s])).logprob();
      loglikelihood += logP[s];
    }
    return loglikelihood;
  }

  protected static double[] uniformdiscrete(int n) {
    double[] ps = new double[n];
    for (int i=0; i<n; i++) ps[i] = 1.0/(n * n);
    return ps;
  }

}
